# -*- coding:utf-8 -*-

from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.metrics import roc_auc_score, accuracy_score, f1_score, precision_score, recall_score, roc_curve, auc, \
    confusion_matrix, precision_recall_curve
import matplotlib.pyplot as plt
import matplotlib as mpl
import catboost as cb
import pandas as pd
import numpy as np
import dataset as ds
import os
import time
import logging
import seaborn as sns

# 配置中文字体支持
mpl.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
mpl.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 创建必要的目录
os.makedirs('./model', exist_ok=True)
os.makedirs('./result', exist_ok=True)


def evaluate_model(model, x_test, y_test):
    """评估模型性能并返回多个指标"""
    start_time = time.time()
    y_pred = model.predict_proba(x_test)
    y_pred_class = model.predict(x_test)
    prediction_time = time.time() - start_time

    auc_score = roc_auc_score(y_test, y_pred[:, 1])
    acc = accuracy_score(y_test, y_pred_class)
    f1 = f1_score(y_test, y_pred_class)
    precision = precision_score(y_test, y_pred_class)
    recall = recall_score(y_test, y_pred_class)

    # 计算混淆矩阵
    cm = confusion_matrix(y_test, y_pred_class)
    tn, fp, fn, tp = cm.ravel()

    logger.info(f"模型评估结果:")
    logger.info(f"  AUC: {auc_score:.6f}")
    logger.info(f"  准确率: {acc:.4f}")
    logger.info(f"  F1分数: {f1:.4f}")
    logger.info(f"  精确率: {precision:.4f}")
    logger.info(f"  召回率: {recall:.4f}")
    logger.info(f"  真阳性(TP): {tp}, 假阳性(FP): {fp}")
    logger.info(f"  假阴性(FN): {fn}, 真阴性(TN): {tn}")
    logger.info(f"  预测耗时: {prediction_time:.4f}秒")

    return {
        'auc': auc_score,
        'accuracy': acc,
        'f1': f1,
        'precision': precision,
        'recall': recall,
        'confusion_matrix': cm
    }


def save_feature_importance(model, feature_names):
    """保存特征重要性图"""
    feature_importance = model.get_feature_importance()
    sorted_idx = np.argsort(feature_importance)[::-1]

    plt.figure(figsize=(12, 8))
    plt.title("特征重要性")
    plt.bar(range(len(feature_importance)), feature_importance[sorted_idx], align='center')
    plt.xticks(range(len(feature_importance)), np.array(feature_names)[sorted_idx], rotation=90)
    plt.tight_layout()

    # 保存特征重要性图
    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
    plt.savefig(f'./result/feature_importance_{timestamp}.png')
    plt.close()

    # 同时保存为CSV
    importance_df = pd.DataFrame({
        'feature': np.array(feature_names)[sorted_idx],
        'importance': feature_importance[sorted_idx]
    })
    importance_df.to_csv(f'./result/feature_importance_{timestamp}.csv', index=False)

    return importance_df


def plot_roc_curve(y_true, y_proba, timestamp):
    """绘制并保存ROC曲线"""
    fpr, tpr, _ = roc_curve(y_true, y_proba)
    roc_auc = auc(fpr, tpr)

    plt.figure(figsize=(10, 8))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC曲线 (AUC = {roc_auc:.4f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('假正率')
    plt.ylabel('真正率')
    plt.title('接收者操作特征曲线')
    plt.legend(loc="lower right")
    plt.grid(True)

    plt.savefig(f'./result/roc_curve_{timestamp}.png')
    plt.close()

    return roc_auc


def plot_precision_recall_curve(y_true, y_proba, timestamp):
    """绘制并保存精确率-召回率曲线"""
    precision, recall, _ = precision_recall_curve(y_true, y_proba)
    average_precision = np.mean(precision)

    plt.figure(figsize=(10, 8))
    plt.plot(recall, precision, color='darkgreen', lw=2, label=f'精确率-召回率曲线 (AP = {average_precision:.4f})')
    plt.xlabel('召回率')
    plt.ylabel('精确率')
    plt.title('精确率-召回率曲线')
    plt.legend(loc="upper right")
    plt.grid(True)

    plt.savefig(f'./result/precision_recall_curve_{timestamp}.png')
    plt.close()

    return average_precision


def plot_confusion_matrix(cm, timestamp):
    """绘制并保存混淆矩阵"""
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=['负类', '正类'],
                yticklabels=['负类', '正类'])
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.title('混淆矩阵')

    plt.savefig(f'./result/confusion_matrix_{timestamp}.png')
    plt.close()


def cross_validate_model(x, y, n_splits=5):
    """执行交叉验证"""
    logger.info(f"开始 {n_splits} 折交叉验证...")
    skf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=42)
    fold_metrics = []

    for fold, (train_idx, val_idx) in enumerate(skf.split(x, y)):
        logger.info(f"训练折叠 {fold + 1}/{n_splits}")
        x_train, x_val = x.iloc[train_idx], x.iloc[val_idx]
        y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

        # 计算类别权重
        scale_pos_weight = len(y_train[y_train == 0]) / len(y_train[y_train == 1])

        model = cb.CatBoostClassifier(
            iterations=2000,
            od_type='Iter',
            # 只使用od_wait或early_stopping_rounds中的一个
            od_wait=100,  # 使用od_wait控制早停
            depth=6,  # 增加深度以捕捉更复杂模式
            learning_rate=0.05,
            l2_leaf_reg=5,
            loss_function='Logloss',
            logging_level='Silent',
            subsample=0.80,
            random_seed=42,
            thread_count=-1,
            eval_metric='AUC',
            scale_pos_weight=scale_pos_weight,  # 使用scale_pos_weight处理类别不平衡
        )

        model.fit(
            x_train, y_train,
            eval_set=(x_val, y_val),
            use_best_model=True,
            verbose=100
        )

        # 评估当前折叠
        y_pred = model.predict_proba(x_val)
        auc_score = roc_auc_score(y_val, y_pred[:, 1])
        y_pred_class = model.predict(x_val)
        f1 = f1_score(y_val, y_pred_class)

        fold_metrics.append({
            'auc': auc_score,
            'f1': f1
        })

        logger.info(f"折叠 {fold + 1} 结果 - AUC: {auc_score:.4f}, F1: {f1:.4f}")

    # 计算平均指标
    avg_auc = np.mean([m['auc'] for m in fold_metrics])
    avg_f1 = np.mean([m['f1'] for m in fold_metrics])

    logger.info(f"交叉验证完成 - 平均AUC: {avg_auc:.4f}, 平均F1: {avg_f1:.4f}")
    return avg_auc, avg_f1


@ds.time_cost
def train():
    try:
        logger.info("加载训练数据...")
        x, y = ds.load_data()
        logger.info(f"数据集形状: x={x.shape}, y={y.shape}")

        # 分析类别分布
        class_counts = y.value_counts()
        imbalance_ratio = class_counts[0] / class_counts[1]
        logger.info(f"类别分布: 负类={class_counts[0]}, 正类={class_counts[1]}, 不平衡比例: {imbalance_ratio:.2f}:1")

        # 保存特征名称
        feature_names = list(x.columns)

        # 执行交叉验证
        avg_auc, avg_f1 = cross_validate_model(x, y, n_splits=3)

        logger.info("划分训练集和测试集...")
        x_train, x_test, y_train, y_test = train_test_split(
            x, y, test_size=0.3, random_state=42, stratify=y
        )
        logger.info(f"训练集: {x_train.shape}, 测试集: {x_test.shape}")
        logger.info(f"训练集类别分布: 负类={sum(y_train == 0)}, 正类={sum(y_train == 1)}")
        logger.info(f"测试集类别分布: 负类={sum(y_test == 0)}, 正类={sum(y_test == 1)}")

        # 计算类别权重
        scale_pos_weight = len(y_train[y_train == 0]) / len(y_train[y_train == 1])
        logger.info(f"类别权重(scale_pos_weight): {scale_pos_weight:.2f}")

        logger.info("训练CatBoost模型...")
        model = cb.CatBoostClassifier(
            iterations=3000,  # 增加迭代次数
            od_type='Iter',
            # 只使用od_wait或early_stopping_rounds中的一个
            od_wait=100,  # 使用od_wait控制早停
            depth=6,  # 增加深度以捕捉更复杂模式
            learning_rate=0.05,
            l2_leaf_reg=5,
            loss_function='Logloss',
            logging_level='Silent',
            subsample=0.80,
            random_seed=42,
            thread_count=-1,
            eval_metric='AUC',
            scale_pos_weight=scale_pos_weight,  # 使用scale_pos_weight处理类别不平衡
        )

        # 训练模型并使用验证集进行早停
        model.fit(
            x_train, y_train,
            eval_set=(x_test, y_test),
            use_best_model=True,
            verbose=100
        )
        logger.info("模型训练完成")

        # 评估模型
        metrics = evaluate_model(model, x_test, y_test)
        timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())

        # 绘制ROC曲线
        y_pred_proba = model.predict_proba(x_test)[:, 1]
        plot_roc_curve(y_test, y_pred_proba, timestamp)

        # 绘制精确率-召回率曲线
        plot_precision_recall_curve(y_test, y_pred_proba, timestamp)

        # 绘制混淆矩阵
        plot_confusion_matrix(metrics['confusion_matrix'], timestamp)

        # 保存特征重要性
        logger.info("保存特征重要性...")
        importance_df = save_feature_importance(model, feature_names)
        top_features = importance_df.head(10)['feature'].tolist()
        logger.info(f"Top 10重要特征: {', '.join(top_features)}")

        # 保存模型
        logger.info("保存模型文件...")
        model_name = f'model_auc_{metrics["auc"]:.4f}_f1_{metrics["f1"]:.4f}_{timestamp}'
        model_path = f'./model/{model_name}.model'
        model.save_model(model_path)
        logger.info(f"模型已保存至: {model_path}")

        return model, metrics

    except Exception as e:
        logger.error(f"训练过程中发生错误: {e}", exc_info=True)
        return None, None


if __name__ == '__main__':
    model, metrics = train()

    if model and metrics:
        logger.info("训练成功完成!")
        logger.info(f"最终AUC分数: {metrics['auc']:.6f}")
        logger.info(f"最终F1分数: {metrics['f1']:.6f}")
    else:
        logger.error("训练失败")